import os, sys
import numpy as np
from scipy.integrate import quad
from scipy.interpolate import interp1d

#   SNeIa are sensitive to very limited number of cosmological parameters:
#   H0 :
#   Omega_m
#   w(z) or w(a)
#   Omega_k


class SNeCosmology:
    def __init__(self, H0=70.0, Omega_m=0.3, Omega_k=0.0, w0=-1.0, wa=0.0, z_sn_max=2.0, grid_size=100):
        """
            H0:         Hubble constant, degenerates with the absolute magnitude of SNe;
            Omega_m:    Fractional matter density w.r.t the critical energy density;
            Omega_k:    curvature contribution to the total energy density;
            w0, wa:     w(a) model parameters, ref CPL;
            z_sn_max:   max redshfit of SNe, specifies the redshift range of the computations;
            grid_size:  the size of the grid that to be used for interpolation.
        """
        self.c_light = 2.99792458e5 # km/sec
        self.H0 = H0
        self.Omega_m = Omega_m
        self.Omega_k = Omega_k
        self.wa = wa
        self.w0 = w0
        self.z_sn_max = z_sn_max
        self.grid_size= grid_size
        self.grid_zmin= 0.005
        self.grid_z  = None
        self.grid_mu = None
        self.mu_interp = None

        self.init_grid()

    def weff(self,z):
        return -1.0

    def iE(self,z):
        E2  = self.Omega_m*(1.0+z)**3 + (1.0-self.Omega_k-self.Omega_m)*(1.0+z)**(3.0*(1.0+self.weff(z)))
        return 1.0/E2**0.5

    def init_grid(self):
        self.grid_z = np.linspace(self.grid_zmin,self.z_sn_max, self.grid_size)
        self.grid_mu= np.zeros(self.grid_size)

        for i in range(self.grid_size):
            zi = self.grid_z[i]
            dc = quad(self.iE,0.0,zi)[0]
            dl = (self.c_light/self.H0)*(1.0+zi)*dc
            # if abs(self.Omega_k) < 1e-5:
            #     dl *= dc
            # elif self.Omega_k > 0:
            #     dl *= np.sinh(np.sqrt(-self.Omega_k)*dc)/np.sqrt(-self.Omega_k)
            # elif self.Omega_k < 0:
            #     dl *= np.sin(np.sqrt(self.Omega_k)*dc)/np.sqrt(self.Omega_k)

            self.grid_mu[i] = 5.0*np.log10(dl) + 25

        self.mu_interp = interp1d(self.grid_z,self.grid_mu,kind='cubic')

    def dist_mu(self,z):
        if self.mu_interp is not None:
            return self.mu_interp(z)
        else:
            print('self.mu is None')
            sys.exit(0)
            
    def dist_mu_ddf(self,z):
        h = 0.01
        return (self.dist_mu(z-h)+self.dist_mu(z+h)-2.0*self.dist_mu(z))/h/h
